#include <cassert>
#include <iostream>

#include "../src/kl.hh"
#include "../src/ordinallogistic.hh"

namespace 
{
  using namespace playerpiano;

  struct LogPTestData
  {
    struct
      {
        float w[17 + 5];
        Feature f[4];
      }                 inputs;
  
    float logp[6];
  };
  
  void
  test_logp (void)
  {
    LogPTestData data[] = 
      {{{{0.641365, 0.541907, 0.716336, 0.59359, 0.353132, 0.313228,  0.408432, 0.590289, 0.446713, 0.772804, 0.347975, 0.90308,  0.364858, 0.214532, 0.655571, 0.38096, 0.72895, 0.0182342,  0.381825, 0.314458, 0.191604,  0.932273}, {{11, 0.0458032}, {16, 0.673896}, {7, 0.550239}, {0,  0.390366}}}, {-2.9806, -2.63729, -2.19452, -1.92618, -1.53509, -0.903547}}, {{{0.329467, 0.080306, 0.197107, 0.0771379, 0.921036,  0.490017, 0.750394, 0.304333, 0.57306, 0.586937, 0.385536,  0.0898011, 0.91749, 0.205978, 0.656586, 0.0715669, 0.535665,  0.891519, 0.464982, 0.139294, 0.489862,  0.217623}, {{1, 0.914744}, {10, 0.748928}, {6, 0.160395}, {0,  0.137317}}}, {-3.64307, -3.15083, -2.40943, -1.78482, -1.04025, -1.13958}}, {{{0.717637, 0.67179, 0.239359, 0.6473, 0.967243,  0.367457, 0.666299, 0.0603631, 0.581708, 0.277656, 0.748809,  0.854385, 0.925122, 0.206089, 0.213144, 0.962866, 0.460139,  0.066795, 0.723282, 0.745243, 0.545395,  0.317867}, {{15, 0.562887}, {16, 0.607926}, {14, 0.827758}, {15,  0.646077}}}, {-4.79742, -4.04557, -3.24788, -2.01775, -1.19557, -0.693331}}, {{{0.323528, 0.960626, 0.860515, 0.27862, 0.657229,  0.900262, 0.278807, 0.000964098, 0.908419, 0.045877, 0.353685,  0.794875, 0.695275, 0.0830109, 0.893546, 0.72808, 0.971993,  0.337768, 0.34815, 0.410214, 0.409106,  0.729842}, {{10, 0.520392}, {4, 0.764137}, {3, 0.0855783}, {6,  0.769217}}}, {-3.31876, -2.53917, -2.39668, -2.15408, -1.32958, -0.883695}}, {{{0.659878, 0.485517, 0.42835, 0.868954, 0.381071,  0.484553, 0.51993, 0.823077, 0.0273855, 0.689678, 0.824655,  0.740066, 0.13384, 0.961597, 0.852662, 0.402298, 0.785689,  0.551384, 0.443556, 0.672456, 0.265297,  0.787247}, {{7, 0.357978}, {11, 0.90324}, {6, 0.60542}, {6,  0.30173}}}, {-4.83167, -4.06262, -3.11921, -2.11863, -1.55866, -0.510799}}, {{{0.929628, 0.0342855, 0.224349, 0.817177, 0.409698,  0.211208, 0.196963, 0.127499, 0.585042, 0.471142, 0.0631236,  0.165902, 0.73238, 0.0688432, 0.277434, 0.614518, 0.288824,  0.396387, 0.0121372, 0.827271, 0.930847,  0.493147}, {{15, 0.406718}, {9, 0.525541}, {0, 0.0012187}, {13,  0.458862}}}, {-1.89552, -2.29492, -1.79899, -1.49311, -1.78007, -1.65955}}, {{{0.182369, 0.708365, 0.591521, 0.247654, 0.985405,  0.580865, 0.00647862, 0.776512, 0.922282, 0.414964, 0.274098,  0.707669, 0.644848, 0.800446, 0.985274, 0.311282, 0.63271,  0.973175, 0.0544271, 0.818134, 0.225993,  0.447634}, {{16, 0.0532084}, {13, 0.359273}, {12, 0.0436239}, {6,  0.739269}}}, {-2.16872, -1.74211, -1.84149, -1.82403, -1.4627, -1.83912}}, {{{0.461687, 0.111619, 0.0582184, 0.158404, 0.455209,  0.335107, 0.135937, 0.74344, 0.181111, 0.627439, 0.491089,  0.942994, 0.195837, 0.316157, 0.858379, 0.969819, 0.14141,  0.498023, 0.632386, 0.522185, 0.0882012,  0.13875}, {{11, 0.588762}, {7, 0.782916}, {2, 0.626514}, {11,  0.0271307}}}, {-3.27715, -3.03264, -2.43807, -1.7901, -1.24231, -0.991324}}, {{{0.530544, 0.624513, 0.171305, 0.692023, 0.394607,  0.881073, 0.990195, 0.0645846, 0.903518, 0.938079, 0.794358,  0.748428, 0.0451393, 0.96826, 0.652948, 0.250405, 0.412753,  0.446075, 0.564747, 0.111655, 0.823991,  0.663158}, {{10, 0.938233}, {11, 0.0845244}, {10, 0.293447}, {15,  0.0386457}}}, {-1.50882, -1.66422, -1.91358, -1.70973, -2.02661, -2.04606}}, {{{0.766928, 0.392501, 0.89884, 0.157573, 0.776734,  0.327916, 0.995322, 0.219494, 0.982376, 0.579489, 0.950183,  0.251234, 0.329428, 0.329084, 0.537429, 0.80516, 0.76468,  0.217429, 0.713439, 0.142001, 0.826447,  0.132904}, {{7, 0.419991}, {6, 0.103356}, {3, 0.0595188}, {14,  0.740403}}}, {-1.17986, -1.61886, -1.68343, -2.25433, -2.08398, -2.53278}}};

    for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
      {
        OrdinalLogisticRegression r (17, 6, 1, data[i].inputs.w, 1);
        std::vector<Feature> features (data[i].inputs.f, data[i].inputs.f + 4);
        IndexedFeatureSet feature_set;
        unsigned int ns = 256 * drand48 ();
        feature_set[ns] = features;

        std::vector<float> logp = r.logp (feature_set);
  
        for (unsigned int j = 0; j < 5; ++j)
          {
            assert (fabs (logp[j] - data[i].logp[j]) <=
                    1e-4 * (1 + fabs (logp[j]) + fabs (data[i].logp[j])) ||
                    (std::cerr << "logp[" << j << "] " << logp[j] 
                               << " ?= " << data[i].logp[j] << std::endl,
                     0));
          }
      }
  }

  struct GradWTestData
  {
    struct
      {
        float w[9 + 3];
        Feature f[2];
        float   q[4];
        unsigned int nitems;
      }                 inputs;
  
    struct
      {
        float logp[4];
        float kl;
        float gradw[9 + 3];
      }                 outputs;
  };

  void
  test_gradw (void)
  {
    GradWTestData data[] = 
    {{{{0.0509479, 0.359804, 0.497142, 0.555551, 0.0446485, 0.739536,  0.622195, 0.116873, 0.0231649, 0.0542884, 0.186905,  0.405299}, {{4, 0.607922}, {8, 0.198295}}, {0.320266, 0.178889,  0.0893629, 0.411482},  4}, {{-1.29163, -1.19116, -1.39954, -1.74527}, -1.46996, {0.0327061, -0.169488, -0.361169, 0, 0, 0, 0, 0.164734, 0, 0, 0,  0.0537336}}}, {{{0.385411, 0.0269173, 0.892468, 0.754985,  0.304649, 0.0985521, 0.84152, 0.395181, 0.807507, 0.543001,  0.796872, 0.655645}, {{6, 0.185311}, {5, 0.426128}}, {0.323145,  0.251162, 0.416993, 0.00869952},  2}, {{-1.54806, -1.48875, -1.07094, -1.51868}, -1.33395, {-0.0822195, 0.122535, -0.235934, 0, 0, 0, 0,  0, -0.194647, -0.0846465, 0, 0}}}, {{{0.165785, 0.403062,  0.475303, 0.728643, 0.0198246, 0.730971, 0.089892, 0.701726,  0.127357, 0.975986, 0.785243,  0.603174}, {{8, 0.285836}, {7, 0.580805}}, {0.398745, 0.0245398,  0.199412, 0.377303},  4}, {{-1.93119, -1.46849, -1.24307, -1.08989}, -1.46519, {0.212323, -0.0527297, -0.159875, 0, 0, 0, 0, 0, 0,  0, -0.151448, -0.0745333}}}, {{{0.792425, 0.634045, 0.715258,  0.323803, 0.794019, 0.613215, 0.549474, 0.920741, 0.318716,  0.884572, 0.529649,  0.189769}, {{4, 0.228824}, {6, 0.182847}}, {0.245401, 0.130409,  0.270587, 0.353604},  5}, {{-0.924992, -1.34499, -1.6066, -1.94943}, -1.52644, {-0.309619, -0.408077, -0.354301, 0, 0, 0, 0, 0.147283, 0, 0.11769, 0,  0}}}, {{{0.116456, 0.632978, 0.465844, 0.519501, 0.627491,  0.707819, 0.67342, 0.885456, 0.912233, 0.384016, 0.879401,  0.272241}, {{3, 0.362759}, {2, 0.463275}}, {0.272844, 0.188649,  0.405413, 0.133095},  6}, {{-1.74268, -1.28693, -1.3477, -1.24134}, -1.42985, {0.0783832, -0.0951717, 0.0782603, 0, 0, -0.122313, -0.095775, 0, 0,  0, 0, 0}}}, {{{0.331861, 0.204822, 0.430818, 0.0597221, 0.888281,  0.625148, 0.314362, 0.426744, 0.422436, 0.105647, 0.686871,  0.718925}, {{3, 0.749017}, {0, 0.220191}}, {0.264642, 0.114416,  0.29709, 0.323852},  5}, {{-1.30362, -1.38687, -1.34308, -1.52529}, -1.39666, {-0.0732759, -0.183307, -0.192457, 0.0562675, 0, 0, 0.191403, 0, 0, 0,  0, 0}}}, {{{0.411878, 0.871635, 0.308931, 0.560282, 0.578768,  0.59813, 0.97707, 0.35546, 0.14795, 0.538407, 0.0887892,  0.730312}, {{3, 0.833588}, {8, 0.111664}}, {0.364033, 0.341258,  0.0801529, 0.214556},  2}, {{-1.98468, -1.50053, -1.47614, -0.88905}, -1.54362, {0.0206691, -0.0909626, 0.0420249, 0, 0, 0, -0.640157, 0, 0, 0,  0, -0.0857524}}}, {{{0.917336, 0.404474, 0.37208, 0.0578288,  0.0477205, 0.456524, 0.960202, 0.186194, 0.73879, 0.896242,  0.381434, 0.588064}, {{1, 0.76172}, {1, 0.540781}}, {0.13217,  0.0281097, 0.380931, 0.458789},  1}, {{-0.667129, -1.52231, -1.86463, -2.17455}, -1.83892, {-1.29835, -0.975581, -0.71721, 0, 1.68967, 0, 0, 0, 0, 0, 0,  0}}}, {{{0.399895, 0.937993, 0.00657746, 0.185805, 0.253177,  0.545255, 0.0892411, 0.781331, 0.881097, 0.487426, 0.0415206,  0.324807}, {{7, 0.920895}, {7, 0.301232}}, {0.152592, 0.216019,  0.271916, 0.359473},  1}, {{-0.832586, -1.18174, -2.06899, -2.02482}, -1.67278, {-0.682226, -1.31105, -0.234033, 0, 0, 0, 0, 0, 0, 0, 1.07895,  0}}}, {{{0.541011, 0.887784, 0.305978, 0.663511, 0.868081,  0.0773147, 0.906083, 0.725518, 0.861503, 0.89151, 0.652906,  0.180263}, {{1, 0.772262}, {6, 0.110179}}, {0.25893, 0.232436,  0.245153, 0.263481},  4}, {{-1.6371, -1.4095, -1.52867, -1.06604}, -1.40715, {-0.0708651, -0.169387, 0.00439452, 0, -0.152781, 0, 0, 0,  0, -0.0217974, 0, 0}}}};

    for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
      {
        OrdinalLogisticRegression rzero (9, 4, data[i].inputs.nitems, data[i].inputs.w, 1);
        std::vector<Feature> features (data[i].inputs.f, data[i].inputs.f + 2);
        IndexedFeatureSet feature_set;
        unsigned int ns = 256 * drand48 ();
        feature_set[ns] = features;
        std::vector<float> logp = rzero.logp (feature_set);

        std::vector<float> w (data[i].inputs.w, data[i].inputs.w + 9 + 3);
        OrdinalLogisticRegression r (9, 4, data[i].inputs.nitems, &w[0], 1);
        std::vector<float> q (data[i].inputs.q, data[i].inputs.q + 4);

        float kl = kldivergence (logp, q);
          
        r.update (1, 1, 0, feature_set, logp, q);

        assert (fabs (kl - data[i].outputs.kl) <=
                1e-4 * (1 + fabs (kl) + fabs (data[i].outputs.kl)) ||
                (std::cerr << "kl = " << kl
                           << " ?= " << data[i].outputs.kl << std::endl,
                 0));

        for (unsigned int j = 0; j < 4; ++j)
          {
            assert (fabs (logp[j] - data[i].outputs.logp[j]) <=
                    1e-4 * (1 + fabs (logp[j]) + fabs (data[i].outputs.logp[j])) ||
                    (std::cerr << "logp[" << j << "] " << logp[j] 
                               << " ?= " << data[i].outputs.logp[j] << std::endl,
                     0));
          }

        for (unsigned int j = 0; j < 9 + 3; ++j)
          {
            w[j] -= data[i].inputs.w[j];

            assert (fabs (w[j] - data[i].outputs.gradw[j]) <=
                    1e-4 * (1 + fabs (w[j]) + fabs (data[i].outputs.gradw[j])) ||
                    (std::cerr << "w[" << j << "] " << w[j] 
                               << " ?= " << data[i].outputs.gradw[j] << std::endl,
                     0));
          }
      }
  }
}

int
main (void)
{
  test_logp ();
  test_gradw ();

  return 0;
}
